"""
标签平滑
@ author： 东阳
"""


import torch
import torch.nn as nn


class LabelSmoothingCrossEntropy(nn.Module):

    def __init__(self, eps=0.01):
        super(LabelSmoothingCrossEntropy, self).__init__()
        self.eps = eps

    def forward(self, x, target):
        # CE(q, p) = - sigma(q_i * log(p_i))
        log_probs = torch.nn.functional.log_softmax(x, dim=-1)  # 实现  log(p_i)
        # H(q, p)
        H_qp = -log_probs.gather(index=target.unsqueeze(1), dim=-1)  # 只需要q_i == 1的地方， 此时已经得到CE
        # print("index:", target.unsqueeze(1), target.unsqueeze(1).shape, )
        # print("log_probs: ", log_probs, log_probs.shape)
        # print("H_qp: ", H_qp, H_qp.shape)

        H_qp = H_qp.squeeze(1)
        # H(u, p)
        H_uq = -log_probs.mean()  # 由于u是均匀分布，等价于求均值

        loss = (1 - self.eps) * H_qp + self.eps * H_uq

        return loss.mean()
